package cc.spider.web.fetcher;

import cc.spider.core.IFetcherListener;
import cc.spider.web.model.GlobalDefaultDataExt;
import cc.spider.web.model.http.ICookie;
import cc.spider.web.model.http.IHeader;
import cc.spider.web.model.http.Request;
import cc.spider.web.model.http.Response;
import okhttp3.*;

import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLSocketFactory;
import javax.net.ssl.TrustManager;
import javax.net.ssl.X509TrustManager;
import java.io.IOException;
import java.security.KeyManagementException;
import java.security.NoSuchAlgorithmException;
import java.security.cert.X509Certificate;
import java.util.*;
import java.util.concurrent.CopyOnWriteArrayList;

public class OkHttpFetcher extends HttpFetcher {
    @Override
    public void httpRequest(Request request, GlobalDefaultDataExt globalConfig, IFetcherListener fetcherListener) {
        try {
            OkHttpClient client = this.createOkHttpClient(request, globalConfig);
            client2Call(request, globalConfig, client).enqueue(new Callback() {
                @Override
                public void onResponse(Call call, okhttp3.Response res) throws IOException {

                    Response response = convert2Response(res, request.isBinaryContent());
                    fetcherListener.success(request, response);
                }

                @Override
                public void onFailure(Call call, IOException e) {
                    fetcherListener.fail(request, null);
                }
            });
        } catch (NoSuchAlgorithmException | KeyManagementException e) {
            fetcherListener.fail(request, null);
        }
    }

    private Response convert2Response(okhttp3.Response res, boolean isBinaryContent) throws IOException {
        ResponseBody body = res.body();
        Response<String> response;
        if (body == null) {
            response = new Response("");
        } else {
            response = new Response(isBinaryContent ? body.bytes() : body.string());
        }
        // okhttp不需要用response来管理cookies
        Map<String, List<String>> headers = res.headers().toMultimap();
        Set<String> names = headers.keySet();
        for (String name : names) {
            List<String> values = headers.get(name);
            for (String value : values) {
                response.addHeader(name, value);
            }
        }
        return response;
    }

    protected CookieJar createCookieJar(Map<String, ICookie> cookiesPool) {
        return new CookieJar() {
            @Override
            public List<okhttp3.Cookie> loadForRequest(HttpUrl httpUrl) {
                if (cookiesPool.size() > 0) {
                    List<ICookie> cookies = new CopyOnWriteArrayList<>(cookiesPool.values());
                    cookies.removeIf(ck -> (ck.getExpiry() != null && (ck.getExpiry().getTime() < System.currentTimeMillis())));
                    return convert2Cookies(cookies);
                } else {
                    return new ArrayList<okhttp3.Cookie>();
                }
            }

            @Override
            public void saveFromResponse(HttpUrl httpUrl, List<okhttp3.Cookie> cookies) {
                for (okhttp3.Cookie cookie : cookies) {
                    ICookie ck = new cc.spider.web.model.http.Cookie(cookie.name(), cookie.value());
                    ck.setDomain(cookie.domain());
                    ck.setPath(cookie.path());
                    ck.setSecure(cookie.secure());
                    ck.setHttpOnly(cookie.httpOnly());
                    ck.setExpiry(new Date(cookie.expiresAt()));
                    cookiesPool.put(cookie.name(), ck);
                }
            }
        };
    }

    protected SSLSocketFactory getSocketFactory(X509TrustManager x509TrustManager) throws NoSuchAlgorithmException, KeyManagementException {
        SSLContext sslContext = SSLContext.getInstance("TLS");
        sslContext.init(null, new TrustManager[]{x509TrustManager}, null);
        return sslContext.getSocketFactory();
    }

    protected X509TrustManager createX509TrustManager() {
        return new X509TrustManager() {
            @Override
            public void checkClientTrusted(X509Certificate[] x509Certificates, String s) {

            }

            @Override
            public void checkServerTrusted(X509Certificate[] x509Certificates, String s) {

            }

            @Override
            public X509Certificate[] getAcceptedIssuers() {
                return new X509Certificate[0];
            }
        };
    }

    protected OkHttpClient createOkHttpClient(Request request, GlobalDefaultDataExt globalConfig) throws NoSuchAlgorithmException, KeyManagementException {
        X509TrustManager x509TrustManager = createX509TrustManager();
        SSLSocketFactory sslSocketFactory = getSocketFactory(x509TrustManager);
        return new OkHttpClient().newBuilder()
                .cookieJar(createCookieJar(globalConfig.getCookiesPool(request.getCookiePoolName())))
                .sslSocketFactory(sslSocketFactory, x509TrustManager)
                .hostnameVerifier((hn, se)->true)
                .build();
    }

    protected Call client2Call(Request request, GlobalDefaultDataExt globalConfig, OkHttpClient client) {
        String bodyStr = request.getRequestBody().getBody();
        RequestBody body = null;
        if (bodyStr != null) {
            body = RequestBody.create(bodyStr, MediaType.parse(request.getRequestBody().getContentType()));
        }
        okhttp3.Request.Builder requestBuilder = new okhttp3.Request.Builder()
                .url(request.getUrl())
                .method(request.getMethod(), body);
        List<IHeader> headers = globalConfig.getHeaders();
        headers.addAll(request.getHeaders());
        for (IHeader header : headers) {
            requestBuilder.addHeader(header.getName(), header.getValue());
        }
        return client.newCall(requestBuilder.build());
    }

    private List<okhttp3.Cookie> convert2Cookies(List<ICookie> cks) {
        List<okhttp3.Cookie> cookies = new CopyOnWriteArrayList<>();
        for (ICookie ck : cks) {
            okhttp3.Cookie.Builder cookieBuilder = new okhttp3.Cookie.Builder();
            cookieBuilder.name(ck.getName())
                    .value(ck.getValue())
                    .domain(ck.getDomain())
                    .path(ck.getPath());
            Date expiryDate = ck.getExpiry();
            if (expiryDate != null) {
                cookieBuilder.expiresAt(expiryDate.getTime());
            }
            if (ck.getHttpOnly()) {
                cookieBuilder.httpOnly();
            }
            if (ck.getSecure()) {
                cookieBuilder.httpOnly();
            }
            cookies.add(cookieBuilder.build());
        }
        return cookies;
    }

}
